 /********************************************************************************************************\
 * Copyright (c) 2006, Dan Kruchinin, <just.asgard@gmail.com>.                                            *
 * All rights reserved.                                                                                   *
 *                                                                                                        *
 * -> safe_alloc.c                                                                                        *
 *  This is realisation of several wrappers on standard libc memory allocation and freing functions,      *
 *  that are safe implemented and support MEM_DEBUG mode, that can help to catch memory leaks and look    *
 *  for memory usage in program, that uses safe_alloc functionality.                                      *
 *                                                                                                        *
 * Redistribution and use of this software in source and binary forms, with or without modification, are  *
 * permitted provided that the following conditions are met:                                              *
 *                                                                                                        *
 * - Redistributions of source code must retain the above                                                 *
 * copyright notice, this list of conditions and the                                                      *
 * following disclaimer.                                                                                  *
 *                                                                                                        *
 * - Redistributions in binary form must reproduce the above                                              *
 * copyright notice, this list of conditions and the                                                      *
 * following disclaimer in the documentation and/or other                                                 *
 * materials provided with the distribution.                                                              *
 *                                                                                                        *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED *
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A *
 * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR *
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT     *
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS    *
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR *
 * TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF   *
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.                                                             *
 \********************************************************************************************************/

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <errno.h>
#include "safe_alloc.h"

static void __mem_error(const char*, ...);

#ifdef MEM_DEBUG
static void __mlist_reg_allocation(void*, size_t, const char*, const int);
static void __mlist_reg_free(void*, const char*, const int);
#endif /* MEM_DEBUG */

void* _safe_malloc(size_t size)
{
  void* ptr;

  if((ptr = malloc(size)) == NULL)
    __mem_error("malloc(): can't allocte memory block with size %zd", size);

  return ptr;
}

void* _safe_calloc(int num, size_t size)
{
  void* ptr;

  if((ptr = calloc(num, size)) == NULL)
    __mem_error("calloc(): can't allocate %d memory blocks with size %zd");

  return ptr;
}

void* _safe_realloc(void* ptr, size_t new_size)
{
  void* nptr;

  if((nptr = realloc(ptr, new_size)) == NULL)
    __mem_error("realloc(): can't reallocate memory block pointed to %p with new size %zd", ptr, new_size);

  return nptr;
}

void  _safe_free(void* ptr)
{
  free(ptr);
}

/*
 * 
 * here starts most interesting part of allocation and freing stratagy,
 * if MEM_DEBUG is enabled, we start to cache information about each allocation
 * and freing in our application. so, after all, we can see full picture of used memory. 
 *
 */
#ifdef MEM_DEBUG

/* summary count of allocations and frees */
static int __sum_allocs = 0, __sum_frees = 0;
static int __df_reg = 0; /* used for indicating registered _dump_memory_dbg with atexit() or not */
/* list, that holds "memory history" */
static struct mlist *__mlist = NULL, *__mlist_last = NULL;

void* _safe_malloc_dbg(size_t size, const char* file, const int line)
{
  void* ptr;

  if(!__df_reg) {
	atexit(_dump_memory_dbg);
	__df_reg = 1;
  }
  if((ptr = malloc(size)) == NULL)
    __mem_error("malloc(): can't allocate memory block with size %zd (%s:%d)", size, file, line);

  __sum_allocs++;
  __mlist_reg_allocation(ptr, size, file, line);

  return ptr;
}

void* _safe_calloc_dbg(int num, size_t size, const char* file, const int line)
{
  void* ptr;

  if(!__df_reg) {
	atexit(_dump_memory_dbg);
	__df_reg = 1;
  }
  if((ptr = calloc(num, size)) == NULL) {
    __mem_error("calloc(): can't allocate %d memory blocks with size %zd (%s:%d)",
		num, size, file, line);
  }

  __sum_allocs++;
  __mlist_reg_allocation(ptr, size, file, line);

  return ptr;
}

void* _safe_realloc_dbg(void* ptr, size_t new_size, const char* file, const int line)
{
  void* nptr;

  if(!__df_reg) {
	atexit(_dump_memory_dbg);
	__df_reg = 1;
  }
  if((nptr = realloc(ptr, new_size)) == NULL) {
    __mem_error("realloc(): can't reallocate memory block pointed to %p with size %zd (%s:%d)",
		ptr, new_size, file, line);
  }

  /* 
   *
   * it can be a several situations, when realloc frees memory, wich pointed by old address:
   * first of it, when pointer is not empty and new size = 0, in this case realoc is equial to free(),
   * the other one, when realloc changes pointer address. in this case it must free previous one. 
   *
   */
  if((ptr != NULL) && ((new_size == 0) || (ptr != nptr))) {
    __sum_frees++;
    __mlist_reg_free(ptr, file, line);
  }
  if((ptr == NULL) || (ptr != nptr)) {
    __sum_allocs++;
    __mlist_reg_allocation(nptr, new_size, file, line);
  }

  return nptr;
}

void _safe_free_dbg(void* ptr, const char* file, const int line)
{
  if(!__df_reg) {
	atexit(_dump_memory_dbg);
	__df_reg = 1;
  }
  
  __sum_frees++;
  __mlist_reg_free(ptr, file, line);
  _safe_free(ptr);
}

/* dump memory allocating picture and free memory history list after it */
void _dump_memory_dbg(void)
{
  struct mlist *tmp, *next = NULL;
  int lost = 0, allocated = 0;

  fprintf(stderr, "\n-------------------------------\n");
  fprintf(stderr, "MEMORY DUMP\n");
  fprintf(stderr, "-------------------------------\n\n");
  
  for(tmp = __mlist; tmp != NULL; tmp = next) {
    fprintf(stderr, "%#x allocated [%d bytes] (%s:%d) ", tmp->addr, tmp->size,
	    tmp->allocated->file, tmp->allocated->line);

    if(tmp->fried != NULL) {
      fprintf(stderr, "[FRIED (%s:%d)]\n", tmp->fried->file, tmp->fried->line);
      _safe_free(tmp->fried->file);
      _safe_free(tmp->fried);
    }
    else {
      lost += tmp->size;
      fprintf(stderr, "[DID NOT FREE]\n");
    }

	allocated += tmp->size;
    next = tmp->next;
    _safe_free(tmp->allocated->file);
    _safe_free(tmp->allocated);
    _safe_free(tmp);
  }

  __mlist = NULL;

  fprintf(stderr, "\n-------------------------------\n");
  fprintf(stderr, "%d allocations, %d frees\n", __sum_allocs, __sum_frees);
  fprintf(stderr, "%d bytes allocated, %d bytes lost...\n", allocated, lost);
  fprintf(stderr, "-------------------------------\n");

  __sum_allocs = __sum_frees = 0;
}

static void __mlist_reg_allocation(void* ptr, size_t size, const char* file, const int line)
{
  struct mlist* tmp;
  size_t len;

  tmp = (struct mlist*)_safe_calloc(1, sizeof(*tmp));
  tmp->allocated = (struct mlist_data*)_safe_calloc(1, sizeof(*(tmp->allocated)));
  tmp->size = size;
  tmp->addr = (uintptr_t)ptr;
  len = strlen(file) + 1;
  tmp->allocated->file = (char*)_safe_malloc(len);
  strncpy(tmp->allocated->file, file, len);
  tmp->allocated->line = line;

  if(__mlist == NULL)
    __mlist = tmp;
  else
    __mlist_last->next = tmp;

  __mlist_last = tmp;
}

static void __mlist_reg_free(void* ptr, const char* file, const int line)
{
  struct mlist* tmp;

  for(tmp = __mlist; tmp != NULL; tmp = tmp->next) {
    if(tmp->addr == (uintptr_t)ptr) {
      size_t len;

      tmp->fried = (struct mlist_data*)_safe_calloc(1, sizeof(*(tmp->fried)));
      len = strlen(file) + 1;
      tmp->fried->file = (char*)_safe_malloc(len);
      strncpy(tmp->fried->file, file, len);
      tmp->fried->line = line;
    }
  }
}

#endif /* MEM_DEBUG */

static void __mem_error(const char* fmt, ...)
{
  va_list ap;
  int errcode = errno;

  va_start(ap, fmt);
  fprintf(stderr, "MEMORY ERROR:\n");
  vfprintf(stderr, fmt, ap);
  va_end(ap);
  fprintf(stderr, "\nERRNO:\n%s\n", strerror(errcode));

  exit(EXIT_FAILURE);
}
